import os
import sys
import argparse
import shutil

WIDTH = 4096
HEIGHT = 2048
# WIDTH = 6400
# HEIGHT = 3200


def main(w, h):
    single_width = int(WIDTH / w)
    single_height = int(HEIGHT / h)
    count = 0
    y = 0
    for i in range(h):
        x = 0
        for j in range(w):
            command = 'ffmpeg -i origin.mp4 -vf crop=' + str(single_width) + ':' + str(single_height) + ':' + str(x) \
                      + ':' + str(y) + ' ./FFmpegCrop/' + str(count) + '.mp4 -y'
            os.system(command)
            count += 1
            x += single_width
        y += single_height


def combineTiles_old(w, h, tileList):
    if not os.path.exists('./combinedTiles'):
        os.mkdir('./combinedTiles')
    single_width = int(WIDTH / w)
    single_height = int(HEIGHT / h)
    count = 0
    command = 'ffmpeg'
    for i in range(h):
        for j in range(w):
            count_str = str(count)
            if tileList[count] == 0:
                tmp = ' -i ./tiles/' + count_str + '/' + count_str

            count += 1
            command += tmp

    command += '-i ./tiles/0/output/audio/en/mp4a/seg-1.m4s -filter_complex \"'
    command += '[0:v][1:v]hstack[t0];[2:v][3:v]hstack[t1];[t0][t1]hstack[h0];'
    command += '[4:v][5:v]hstack[t2];[6:v][7:v]hstack[t3];[t2][t3]hstack[h1];'
    command += '[8:v][9:v]hstack[t4];[10:v][11:v]hstack[t5];[t4][t5]hstack[h2];'
    command += '[12:v][13:v]hstack[t6];[14:v][15:v]hstack[t7];[t6][t7]hstack[h3];'
    command += '[h0][h1]vstack[t8];[h2][h3]vstack[t9];[t8][t9]vstack[v]\"'
    command += ' -map \"[v]\" -map 16:a -c:a copy -shortest ./combinedTiles/out.mp4'
    os.system(command)


def combineTiles(w, h, tileList_0, tileList_1, tileList_2):
    if os.path.exists('./combinedTiles'):
        shutil.rmtree('./combinedTiles')
    os.mkdir('./combinedTiles')
    os.mkdir('./combinedTiles/BD')
    os.mkdir('./combinedTiles/HD')
    os.mkdir('./combinedTiles/SD')
    for a in range(len(tileList_0)):
        count = 0
        command = 'ffmpeg'
        for i in range(h):
            for j in range(w):
                count_str = str(count)
                tmp = ''
                if tileList_0[a][count] == 0:
                    tmp = ' -i ./tiles/' + count_str + '/BD/seg_%03d.mp4' % (a + 1)
                elif tileList_0[a][count] == 1:
                    tmp = ' -i ./tiles/' + count_str + '/HD/seg_%03d.mp4' % (a + 1)
                elif tileList_0[a][count] == 2:
                    tmp = ' -i ./tiles/' + count_str + '/SD/seg_%03d.mp4' % (a + 1)
                count += 1
                command += tmp
        command += ' -i origin.m4a -filter_complex \"'
        command += '[0:v][1:v]hstack[t0];[2:v][3:v]hstack[t1];[t0][t1]hstack[h0];'
        command += '[4:v][5:v]hstack[t2];[6:v][7:v]hstack[t3];[t2][t3]hstack[h1];'
        command += '[8:v][9:v]hstack[t4];[10:v][11:v]hstack[t5];[t4][t5]hstack[h2];'
        command += '[12:v][13:v]hstack[t6];[14:v][15:v]hstack[t7];[t6][t7]hstack[h3];'
        command += '[h0][h1]vstack[t8];[h2][h3]vstack[t9];[t8][t9]vstack[v]\"'
        command += ' -map \"[v]\" -map 16:a -c:a copy -shortest -y ' \
                   './combinedTiles/BD/seg_%03d.mp4' % (a + 1)
        os.system(command)

    for a in range(len(tileList_1)):
        count = 0
        command = 'ffmpeg -c:v h264_cuvid'
        for i in range(h):
            for j in range(w):
                count_str = str(count)
                tmp = ''
                if tileList_1[a][count] == 0:
                    tmp = ' -i ./tiles/' + count_str + '/BD/seg_%03d.mp4' % (a + 1)
                elif tileList_1[a][count] == 1:
                    tmp = ' -i ./tiles/' + count_str + '/HD/seg_%03d.mp4' % (a + 1)
                elif tileList_1[a][count] == 2:
                    tmp = ' -i ./tiles/' + count_str + '/SD/seg_%03d.mp4' % (a + 1)
                count += 1
                command += tmp
        command += ' -i origin.m4a -filter_complex \"'
        command += '[0:v][1:v]hstack[t0];[2:v][3:v]hstack[t1];[t0][t1]hstack[h0];'
        command += '[4:v][5:v]hstack[t2];[6:v][7:v]hstack[t3];[t2][t3]hstack[h1];'
        command += '[8:v][9:v]hstack[t4];[10:v][11:v]hstack[t5];[t4][t5]hstack[h2];'
        command += '[12:v][13:v]hstack[t6];[14:v][15:v]hstack[t7];[t6][t7]hstack[h3];'
        command += '[h0][h1]vstack[t8];[h2][h3]vstack[t9];[t8][t9]vstack[v]\"'
        command += ' -map \"[v]\" -map 16:a -c:a copy -shortest -y -c:v h264_nvenc ' \
                   './combinedTiles/HD/seg_%03d.mp4' % (a + 1)
        os.system(command)

    for a in range(len(tileList_2)):
        count = 0
        command = 'ffmpeg -c:v h264_cuvid'
        for i in range(h):
            for j in range(w):
                count_str = str(count)
                tmp = ''
                if tileList_2[a][count] == 0:
                    tmp = ' -i ./tiles/' + count_str + '/BD/seg_%03d.mp4' % (a + 1)
                elif tileList_2[a][count] == 1:
                    tmp = ' -i ./tiles/' + count_str + '/HD/seg_%03d.mp4' % (a + 1)
                elif tileList_2[a][count] == 2:
                    tmp = ' -i ./tiles/' + count_str + '/SD/seg_%03d.mp4' % (a + 1)
                count += 1
                command += tmp
        command += ' -i origin.m4a -filter_complex \"'
        command += '[0:v][1:v]hstack[t0];[2:v][3:v]hstack[t1];[t0][t1]hstack[h0];'
        command += '[4:v][5:v]hstack[t2];[6:v][7:v]hstack[t3];[t2][t3]hstack[h1];'
        command += '[8:v][9:v]hstack[t4];[10:v][11:v]hstack[t5];[t4][t5]hstack[h2];'
        command += '[12:v][13:v]hstack[t6];[14:v][15:v]hstack[t7];[t6][t7]hstack[h3];'
        command += '[h0][h1]vstack[t8];[h2][h3]vstack[t9];[t8][t9]vstack[v]\"'
        command += ' -map \"[v]\" -map 16:a -c:a copy -shortest -y -c:v h264_nvenc ' \
                   './combinedTiles/SD/seg_%03d.mp4' % (a + 1)
        os.system(command)

    file = open('./combinedTiles/BD/segment_list.txt', 'w')
    for i in range(len(tileList_0)):
        file.write('file \'seg_%03d.mp4\'\n' % (i + 1))
    file.close()

    file = open('./combinedTiles/HD/segment_list.txt', 'w')
    for i in range(len(tileList_0)):
        file.write('file \'seg_%03d.mp4\'\n' % (i + 1))
    file.close()

    file = open('./combinedTiles/SD/segment_list.txt', 'w')
    for i in range(len(tileList_0)):
        file.write('file \'seg_%03d.mp4\'\n' % (i + 1))
    file.close()


def splitTile(w, h):
    if os.path.exists('./tiles'):
        shutil.rmtree('./tiles')
    os.mkdir('./tiles')

    segments = os.listdir('./segment')
    segment_count = len(segments)

    # 1024*512 -> 4096*2048
    single_width = int(WIDTH / w)
    single_height = int(HEIGHT / h)
    count = 0
    y = 0
    for i in range(h):
        x = 0
        for j in range(w):
            path = './tiles/' + str(count) + '/'
            if not os.path.exists(path):
                os.mkdir(path)
                os.mkdir(path + '/BD/')

            for k in range(segment_count):
                os.system('ffmpeg -i ./segment/origin_%03d.mp4 -vf crop=' % (k + 1) + str(single_width) + ':'
                          + str(single_height) + ':' + str(x) + ':' + str(y) + ' ' + path
                          + 'BD/seg_%03d.mp4 -y' % (k + 1))

            count += 1
            x += single_width
        y += single_height


# Function: Transcoding the un-segmented videos
def transcoding_old(w, h):
    count = 0
    for i in range(h):
        for j in range(w):
            count_str = str(count)
            path = './tiles/' + count_str + '/'
            #
            command_1 = 'ffmpeg -i ' + path + '1000k.mp4 -b:v 500k ' + path + '500k.mp4 -y'
            os.system(command_1)

            #
            command_2 = 'ffmpeg -i ' + path + '1000k.mp4 -b:v 250k ' + path + '250k.mp4 -y'
            os.system(command_2)
            count += 1


def transcoding(w, h):
    segments = os.listdir('./segment')
    segment_count = len(segments)
    os.chdir('./tiles/')
    count = 0
    for i in range(h):
        for j in range(w):
            count_str = str(count)
            os.chdir('./' + count_str)

            os.mkdir('./SD')
            os.mkdir('./HD')

            for k in range(segment_count):
                os.system('ffmpeg -hwaccel cuvid -c:v h264_cuvid -i ./BD/seg_%03d.mp4 -b:v 750k -c:v h264_nvenc '
                          './HD/seg_%03d.mp4 -y' % (k + 1, k + 1))
                os.system('ffmpeg -hwaccel cuvid -c:v h264_cuvid -i ./BD/seg_%03d.mp4 -b:v 250k -c:v h264_nvenc '
                          './SD/seg_%03d.mp4 -y' % (k + 1, k + 1))
            count += 1
            os.chdir('..')


# May have problems when combining tiles
# Function: split the tiled video into segments
def splitSegment_old(w, h):
    count = 0
    os.chdir('./tiles/')
    for i in range(h):
        for j in range(w):
            count_str = str(count)
            os.chdir('./' + count_str)
            print(os.getcwd())
            if os.path.exists('./output'):
                shutil.rmtree('./output')
            if os.path.isfile('./frag_' + count_str + '_1000k.mp4'):
                os.remove('./frag_' + count_str + '_1000k.mp4')
                os.remove('./frag_' + count_str + '_500k.mp4')
                os.remove('./frag_' + count_str + '_250k.mp4')
            command_1_frag = 'mp4fragment ' + count_str + '_1000k.mp4 ' \
                             + 'frag_' + count_str + '_1000k.mp4'
            os.system(command_1_frag)

            command_2_frag = 'mp4fragment ' + count_str + '_500k.mp4 ' \
                             + 'frag_' + count_str + '_500k.mp4'
            os.system(command_2_frag)

            command_3_frag = 'mp4fragment ' + count_str + '_250k.mp4 ' \
                             + 'frag_' + count_str + '_250k.mp4'
            os.system(command_3_frag)

            command_dash = 'mp4dash frag_' + count_str + '_1000k.mp4 frag_' + count_str + \
                           '_500k.mp4 frag_' + count_str + '_250k.mp4'
            os.system(command_dash)
            os.chdir('..')
            count += 1
    os.chdir('..')


def splitSegment(s):
    if not os.path.exists('./segments'):
        os.mkdir('./segment')
        os.system('mp4box -split ' + str(s) + ' origin.mp4 -out ./segment/')


def mp4BoxDash_old(w, h):
    os.chdir('./tiles/')
    count = 0
    for i in range(h):
        for j in range(w):
            os.chdir('./' + str(count))
            command = 'MP4Box -dash 2000 -frag 2000 -rap -frag-rap -profile onDemand -bs-switching no 1000k.mp4' \
                      ' 500k.mp4 250k.mp4 ../../origin.m4a -out out.mpd'
            os.system(command)
            count += 1
            os.chdir('..')
    os.chdir('..')


def ffmpegCutSegment_old(w, h):
    os.chdir('./tiles/')
    count = 0
    for i in range(h):
        for j in range(w):
            os.chdir('./' + str(count))
            if os.path.exists('./output'):
                shutil.rmtree('./output')
            os.mkdir('./output')
            os.mkdir('./output/0')
            os.mkdir('./output/1')
            os.mkdir('./output/2')
            # ffmpeg -re -i 250k.mp4 -c copy -f segment -segment_time 2 -segment_format mp4 ./output/2/250k-%d.mp4
            os.system('ffmpeg -re -i 250k.mp4 -c copy -f segment -segment_time 2 -segment_format mp4 '
                      './output/2/250k-%d.mp4')
            os.system('ffmpeg -re -i 500k.mp4 -c copy -f segment -segment_time 2 -segment_format mp4 '
                      './output/1/500k-%d.mp4')
            os.system('ffmpeg -re -i 1000k.mp4 -c copy -f segment -segment_time 2 -segment_format mp4 '
                      './output/0/1000k-%d.mp4')
            os.chdir('..')
            count += 1
    os.chdir('..')


def combineSegment():
    os.chdir('./combinedTiles/BD')
    os.system('ffmpeg -f concat -c:v h264_cuvid -i segment_list.txt -y -c:v h264_nvenc ../BD.mp4')
    os.chdir('../HD')
    os.system('ffmpeg -f concat -c:v h264_cuvid -i segment_list.txt -y -c:v h264_nvenc ../HD.mp4')
    os.chdir('../SD')
    os.system('ffmpeg -f concat -c:v h264_cuvid -i segment_list.txt -y -c:v h264_nvenc ../SD.mp4')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process video')
    parser.add_argument('-a', '--getAudio', action='store_true', help='get audio from a video')
    parser.add_argument('-ss', '--splitSegment', action='store_true',
                        help='split origin video into segments with mp4box')
    parser.add_argument('-sd', '--segmentDuration', type=int, default=3, help='segment duration')

    parser.add_argument('-st', '--splitTile', action='store_true', help='split the video into tiles')

    parser.add_argument('-tc', '--transcoding', action='store_true', help='transcode the video into different bitrates')

    parser.add_argument('-ct', '--combineTiles', action='store_true',
                        help='combine tiles with different bitrates according to the given list')

    parser.add_argument('-cs', '--combineSegments', action='store_true',
                        help='combine segments to a single video')
    # parser.add_argument('-p', '--process', action='store_true', help='process the video')
    parser.add_argument('-wt', '--width', type=int, default=4, help='number of tile in horizontal direction')
    parser.add_argument('-ht', '--height', type=int, default=4, help='number of tile in vertical direction')

    args = parser.parse_args()

    # 默认为4*4，如果可变->combineTiles需要更改
    w = args.width
    h = args.height
    s = args.segmentDuration
    # 0 最高清 2 最低清
    tileList_0 = [
        [2, 2, 2, 2,
         2, 0, 0, 2,
         2, 0, 0, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 0, 0,
         2, 2, 0, 0,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         0, 0, 2, 2,
         0, 0, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 0, 0, 2,
         2, 0, 0, 2,
         2, 2, 2, 2],

        [2, 0, 0, 2,
         2, 0, 0, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 0, 0, 2,
         2, 0, 0, 2],

        [2, 2, 2, 2,
         2, 0, 0, 2,
         2, 0, 0, 2,
         2, 2, 2, 2],

        [2, 0, 0, 2,
         2, 0, 0, 2,
         2, 0, 0, 2,
         2, 0, 0, 2],

        [2, 2, 2, 2,
         0, 0, 0, 0,
         0, 0, 0, 0,
         2, 2, 2, 2],

        [0, 0, 2, 2,
         0, 0, 2, 2,
         2, 0, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 0, 0, 2,
         2, 0, 0, 2,
         2, 2, 2, 2]
    ]

    tileList_1 = [
        [2, 2, 2, 2,
         2, 1, 1, 2,
         2, 1, 1, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 1, 1,
         2, 2, 1, 1,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         1, 1, 2, 2,
         1, 1, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 1, 1, 2,
         2, 1, 1, 2,
         2, 2, 2, 2],

        [2, 1, 1, 2,
         2, 1, 1, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 1, 1, 2,
         2, 1, 1, 2],

        [2, 2, 2, 2,
         2, 1, 1, 2,
         2, 1, 1, 2,
         2, 2, 2, 2],

        [2, 1, 1, 2,
         2, 1, 1, 2,
         2, 1, 1, 2,
         2, 1, 1, 2],

        [2, 2, 2, 2,
         1, 1, 1, 1,
         1, 1, 1, 1,
         2, 2, 2, 2],

        [1, 1, 2, 2,
         1, 1, 2, 2,
         2, 1, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 1, 1, 2,
         2, 1, 1, 2,
         2, 2, 2, 2]
    ]

    tileList_2 = [
        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2],

        [2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2,
         2, 2, 2, 2]
    ]

    if args.getAudio:
        os.system('ffmpeg -i origin.mp4 -vn -codec copy origin.m4a')

    elif args.splitSegment:
        splitSegment(s)

    elif args.splitTile:
        splitTile(w, h)

    elif args.transcoding:
        transcoding(w, h)

    elif args.combineTiles:
        combineTiles(w, h, tileList_0, tileList_1, tileList_2)

    elif args.combineSegments:
        combineSegment()

# 分tile、每个tile一个文件夹、转码、mp4dash(分片)
